from backbones.image_classifier_model import ImageClassifierModel
from dataset.dataloader import generate_cars_loader
import torch
from torch import nn
from utils.train_utils import train

if __name__ == '__main__':
    device = torch.device("cuda")
    train_loader, valid_loader = generate_cars_loader(batch_size=2)
    model = ImageClassifierModel(1024, 8, 2, num_classes=50)
    model = model.to(device)
    model.device = device

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    train(train_loader, valid_loader, model, criterion, optimizer, epochs=1000)
